/*
  Copyright (C) 2017 Open Intelligent Machines Co.,Ltd

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/
#include "crowd_count_predict.h"

#include "mxnet/c_predict_api.h"
#include <sys/time.h>
unsigned long get_cur_time(void) {
  struct timeval tv;
  unsigned long ts;

  gettimeofday(&tv, NULL);

  ts = tv.tv_sec * 1000000 + tv.tv_usec;
  return ts;
}

CrowdCountPredictor::CrowdCountPredictor(const std::string &param_file,
                                         const std::string &json_file,
                                         int batch, int channel, int input_h,
                                         int input_w, string input_name,
                                         int ngpuid) {
  Net_ = NULL;
  ngpuid_ = ngpuid;
  num_channels_ = channel;
  input_h_ = input_h;
  input_w_ = input_w;
  input_name_ = input_name;
  Net_ = init(param_file, json_file, batch, num_channels_, input_h_, input_w_);
  cout << "Net_:" << Net_ << endl;
}

CrowdCountPredictor::~CrowdCountPredictor() {
  int res = -2;
  if (NULL != Net_)
    res = MXPredFree(Net_);
}

int CrowdCountPredictor::LoadFile(const std::string &fname,
                                  std::vector<char> &buf) {
  std::ifstream fs(fname, std::ios::binary | std::ios::in);
  if (!fs.good()) {
    std::cerr << fname << " does not exist" << std::endl;
    return -1;
  }
  fs.seekg(0, std::ios::end);
  int fsize = fs.tellg();
  fs.seekg(0, std::ios::beg);
  buf.resize(fsize);
  fs.read(buf.data(), fsize);
  fs.close();
  return 0;
}

void *CrowdCountPredictor::init(const std::string &param_file,
                                const std::string &json_file, int batch,
                                int channel, int input_h, int input_w) {
  std::vector<char> param_buffer;
  std::vector<char> json_buffer;
  PredictorHandle pred_hnd;

  if (LoadFile(param_file, param_buffer) < 0)
    return nullptr;

  if (LoadFile(json_file, json_buffer) < 0)
    return nullptr;

  int device_type = 2;
  mx_uint num_input_nodes = 1;
  const char *input_keys[1];
  const mx_uint input_shape_indptr[] = {0, 4};
  const mx_uint input_shape_data[] = {
      static_cast<mx_uint>(batch), static_cast<mx_uint>(channel),
      static_cast<mx_uint>(input_h), static_cast<mx_uint>(input_w)};

  input_keys[0] = input_name_.c_str();

  MXPredCreate(json_buffer.data(), param_buffer.data(), param_buffer.size(),
               device_type, ngpuid_, num_input_nodes, input_keys,
               input_shape_indptr, input_shape_data, &pred_hnd);

  return pred_hnd;
}
void CrowdCountPredictor::set_input_buffer(std::vector<cv::Mat> &input_channels,
                                           float *input_data, const int height,
                                           const int width) {
  for (int i = 0; i < num_channels_; ++i) {
    cv::Mat channel(height, width, CV_32FC1, input_data);
    input_channels.push_back(channel);
    input_data += width * height;
  }
}

size_t CrowdCountPredictor::run(cv::Mat img, cv::Rect rect) {
  std::vector<float> input(input_h_ * input_w_);

  std::vector<cv::Mat> input_channels;
  double time0 = get_cur_time();

  set_input_buffer(input_channels, input.data(), input_h_, input_w_);

  cv::Mat sample = img.clone();

  cv::Mat sample_resized;
  cv::Size input_geometry_ = Size(input_w_, input_h_);
  if (sample.size() != input_geometry_)
    cv::resize(sample, sample_resized, input_geometry_);
  else
    sample_resized = sample;

  cv::Mat sample_float;
  if (num_channels_ == 3)
    sample_resized.convertTo(sample_float, CV_32FC3);
  else
    sample_resized.convertTo(sample_float, CV_32FC1);

  cv::split(sample_float, input_channels);

  MXPredSetInput(Net_, input_name_.c_str(), input.data(), input.size());
  double time1 = get_cur_time();
  // cout<<"time 1:"<<time1-time0<<endl;
  MXPredForward(Net_);
  double time2 = get_cur_time();
  // cout<<"time 2:"<<time2-time1<<endl;

  mx_uint *shape = NULL;
  mx_uint shape_len = 0;

  MXPredGetOutputShape(Net_, 0, &shape, &shape_len);
  // cout<<"shape_len:"<<shape_len<<endl;

  int reg_size = 1;

  for (unsigned int i = 0; i < shape_len; i++) {
    reg_size *= shape[i];
    // cout<<"shape :"<<i<<" "<<shape[i]<<endl;
  }
  std::vector<float> vfeature(reg_size);
  double time3 = get_cur_time();
  // cout<<"time 3:"<<time3-time2<<endl;

  MXPredGetOutput(Net_, 0, vfeature.data(), reg_size);
  double time4 = get_cur_time();
  // cout<<"time 4:"<<time4-time3<<endl;

  size_t output_w = shape[3];
  size_t output_h = shape[2];
  cv::Mat ma = cv::Mat(output_h, output_w, CV_32FC1);
  for (size_t i = 0; i < vfeature.size(); i++) {
    ma.at<float>(i / output_w, i - i / output_w * output_w) = vfeature[i];
  }
  size_t scale_x = input_w_ / output_w;
  size_t scale_y = input_h_ / output_h;
  Rect rect_resized = Rect(rect.x / scale_x, rect.y / scale_y,
                           rect.width / scale_x, rect.height / scale_y);
  // cout << "rect_resized : "<<rect_resized<<endl;
  Mat img_roi = ma(rect_resized);
  Scalar s_sum = sum(img_roi);
  return size_t(s_sum.val[0]);

  /*//压缩行，成1行*output_w列
      Mat mb(1, output_w, CV_32FC1, Scalar(0));
  reduce(ma, mb, 0, CV_REDUCE_SUM);
  //压缩列，成1行*1列
      Mat mc(1, 1, CV_32FC1, Scalar(0));
  reduce(mb, mc, 1, CV_REDUCE_SUM);

      //cout << "行向量" << endl;
      //cout << mc << endl;

  return size_t(mc.at<float>(0,0));*/
}
